import os
from mmocr.apis import MMOCRInferencer

class MMOCRExtractor:
    """
    一个使用商汤科技MMOCR工具箱进行文字提取的封装类。
    """

    def __init__(self, det_model='DBNet', rec_model='SAR', device='cpu'):
        """
        初始化MMOCR推理器。

        Args:
            det_model (str): 使用的文本检测模型。默认为 'DBNet'。
            rec_model (str): 使用的文本识别模型。默认为 'SAR'，对中文有较好支持。
            device (str): 推理设备，'cuda:0' 或 'cpu'。
        """
        print(f"正在初始化MMOCR推理器，使用设备: {device}...")
        print(f"检测模型: {det_model}, 识别模型: {rec_model}")
        
        # MMOCR会自动下载并缓存模型
        self.inferencer = MMOCRInferencer(det=det_model, rec=rec_model, device=device)
        print("MMOCR推理器初始化完成。")

    def extract_text(self, image_path: str) -> dict:
        """
        从单张图片中提取文字及其位置信息。

        Args:
            image_path (str): 图片文件的路径。

        Returns:
            dict: 包含提取结果的字典，包括识别出的文本和每个文本框的位置。
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"错误: 图片文件未找到 - {image_path}")

        print(f"\n正在使用MMOCR从图片中提取文字: {image_path}")
        # 执行推理
        result = self.inferencer(image_path, return_vis=False)
        
        return result
